'''
使用组合的方式封装自定义的线程池（使用反射封装的做法）
'''
import threading
import queue
import atexit
import logging
from concurrent.futures import ThreadPoolExecutor, Future

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


class MyBoundedThreadPool:
    def __init__(self, max_workers=4, queue_size=100, thread_name_prefix='my-worker'):
        self.executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix=thread_name_prefix)
        self.task_queue = queue.Queue(maxsize=queue_size)
        self._shutdown = threading.Event()

        # 启动任务分发线程
        self._worker_thread = threading.Thread(target=self._dispatch, daemon=True)
        self._worker_thread.start()

        # 启动监控线程
        self._monitor_thread = threading.Thread(target=self._monitor, daemon=True)
        self._monitor_thread.start()

        atexit.register(self.shutdown)

    def _dispatch(self):
        while not self._shutdown.is_set():
            try:
                func, args, kwargs, future = self.task_queue.get(timeout=0.5)
            except queue.Empty:
                continue

            if future.set_running_or_notify_cancel():
                try:
                    result = func(*args, **kwargs)
                    future.set_result(result)
                except Exception as e:
                    future.set_exception(e)
            self.task_queue.task_done()

    def _monitor(self):
        while not self._shutdown.is_set():
            logger.info(f"[ThreadPool] Queue size: {self.task_queue.qsize()}/{self.task_queue.maxsize}")
            time.sleep(1)

    def submit(self, fn, *args, **kwargs) -> Future:
        future = Future()
        try:
            self.task_queue.put((fn, args, kwargs, future), block=True)
        except Exception as e:
            future.set_exception(e)
        return future

    def shutdown(self, wait=True):
        logger.info("Shutting down MyBoundedThreadPool...")
        self._shutdown.set()
        if wait:
            self.task_queue.join()
        self.executor.shutdown(wait=wait)
        logger.info("Shutdown complete.")


if __name__ == '__main__':
    import time


    def task1():
        time.sleep(5)
        return 1


    def task2():
        time.sleep(8)
        return 2


    pool = MyBoundedThreadPool(max_workers=1, queue_size=5)

    # 提交任务
    future1 = pool.submit(task1)
    future2 = pool.submit(task2)

    # 等待结果
    result1 = future1.result()
    result2 = future2.result()

    total = result1 + result2
    print(f"Total result = {total}")

    pool.shutdown()
